import torch
from torch.utils.data import DataLoader
from torch import nn, optim
from datasets.dataset import TicketDataset
from models.price_predictor import PricePredictor
import os
import json
from tqdm import tqdm

CONFIG_PATH = "config/model_config.json"
DATA_PATH = "data/processed/train.pkl"
CHECKPOINT_PATH = "checkpoints/best_model.pt"
BATCH_SIZE = 32
NUM_EPOCHS = 20
LEARNING_RATE = 1e-3
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

os.makedirs("checkpoints", exist_ok=True)

# 加载数据
dataset = TicketDataset(DATA_PATH)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

# 加载模型结构配置
with open(CONFIG_PATH, "r") as f:
    model_cfg = json.load(f)["model"]

input_dim = dataset[0][0].shape[0]
model = PricePredictor(CONFIG_PATH, input_dim).to(DEVICE)

# 训练配置
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# 训练循环
best_loss = float("inf")
for epoch in range(1, NUM_EPOCHS + 1):
    model.train()
    total_loss = 0.0
    loop = tqdm(dataloader, desc=f"Epoch {epoch}")
    for x, y in loop:
        x, y = x.to(DEVICE), y.to(DEVICE)
        optimizer.zero_grad()
        pred = model(x)
        loss = criterion(pred, y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * x.size(0)
        loop.set_postfix(loss=loss.item())

    avg_loss = total_loss / len(dataloader.dataset)
    print(f"✅ Epoch {epoch} finished: Avg Train Loss = {avg_loss:.4f}")

    if avg_loss < best_loss:
        best_loss = avg_loss
        torch.save(model.state_dict(), CHECKPOINT_PATH)
        print("🎯 最佳模型已保存")